#include <Python.h>

#include <numpy/arrayobject.h>
#include <kuai/sim/energy.hpp>
#include <kuai/sim/nonbond.hpp>

#include "pytool.hpp"
#include "EnergyProxy.hpp"


namespace kuai { namespace suan {
	
	EnergyProxy::~EnergyProxy()
	{ }

	bool ForceFieldEnergyProxy::prepare(PyObject* self, PyObject* model) 
	{
		PBC* pbc;
		if (!getPBC(model, &pbc)) {
			return false;
		};

		PyObject* name = PyObject_GetAttrString(self, "name");
		if (name == NULL) {
			return false;
		}

		const char* szName = PyString_AsString(name);
		
		if (szName == NULL) {
			Py_DECREF(name);
			return false;
		}
		else {
			Py_DECREF(name);
		}

		BasicForceFieldEnergy* pfunc = createValenceEnergyFunction(szName, pbc);
		func = std::auto_ptr<EnergyFunction>(pfunc);
		pfunc->pbc = pbc;

		boost::tuple<PyObject*, RealNumber*, size_t> par = getRealArray(model, "parameters");
		if (boost::get<0>(par) == NULL)  {
			return false;
		}
		pfunc->pars = boost::get<1>(par);
		Py_DECREF(boost::get<0>(par));

		boost::tuple<PyObject*, XYZ*, size_t> coords = getXYZArray(model, "coords");
		if (boost::get<0>(coords) == NULL) {
			return false;
		}
		pfunc->coords = boost::get<1>(coords);
		Py_DECREF(boost::get<0>(coords));
		
		boost::tuple<PyObject*, Index*, size_t> lists = getIndexArray(self, "lists");
		if (boost::get<0>(lists) == NULL) {
			return false;
		}
		pfunc->lists = boost::get<1>(lists);
		pfunc->nlist = boost::get<2>(lists);
		Py_DECREF(boost::get<0>(lists));
		
		return true;
	}

	bool NonbondEnergyProxy::prepare(PyObject* self, PyObject* model) {
		PBC* pbc;
		if (!getPBC(model, &pbc)) {
			return false;
		};

		PyObject* name = PyObject_GetAttrString(self, "name");
		if (name == NULL) {
			return false;
		}

		const char* szName = PyString_AsString(name);
		
		if (szName == NULL) {
			Py_DECREF(name);
			return false;
		}
		else {
			Py_DECREF(name);
		}

		bool hasCharge = (PyObject_HasAttrString(model, "charges") == 1);

		hasCharge = hasCharge && (PyObject_HasAttrString(self, "columnb_constant") == 1);

		BasicNonbondEnergy* pfunc = createNonbondEnergyFunction(szName, pbc, hasCharge);
		func = std::auto_ptr<EnergyFunction>(pfunc);
		pfunc->pbc = pbc;

		boost::tuple<PyObject*, RealNumber*, size_t> par = getRealArray(model, "parameters");
		if (boost::get<0>(par) == NULL)  {
			return false;
		}
		pfunc->pars = boost::get<1>(par);
		Py_DECREF(boost::get<0>(par));

		boost::tuple<PyObject*, XYZ*, size_t> coords = getXYZArray(model, "coords");
		if (boost::get<0>(coords) == NULL) {
			return false;
		}
		pfunc->coords = boost::get<1>(coords);
		pfunc->nAtoms = boost::get<2>(coords);
		Py_DECREF(boost::get<0>(coords));
		
		if (PyObject_HasAttrString(self, "lists")) {
			boost::tuple<PyObject*, Index*, size_t> lists = getIndexArray(self, "lists");
			if (boost::get<0>(lists) == NULL) {
				return false;
			}
			pfunc->lists = boost::get<1>(lists);
			pfunc->nlist = boost::get<2>(lists);
			Py_DECREF(boost::get<0>(lists));
		}
		else {
			assert (PyObject_HasAttrString(self, "filter"));
			pfunc->lists = NULL;
			pfunc->nlist = 0;

			PyObject* filter = PyObject_GetAttrString(self, "filter");
			Py_DECREF(filter);
			pfunc->filter =(kuai::Byte*)(PyArray_DATA(filter));

			PyObject* par0 = PyObject_GetAttrString(self, "par0");
			if (par0 == NULL) {
				return false;
			}
			Py_ssize_t pp = PyInt_AsSsize_t(par0);
			if (pp == -1) {
				return false;
			}
			pfunc->par0 = pp;

			boost::tuple<PyObject*, Index*, size_t> atomtype = getIndexArray(model, "atomtype");
			if (boost::get<0>(atomtype) == NULL) {
				return false;
			}
			pfunc->atomType = boost::get<1>(atomtype);
			Py_DECREF(boost::get<0>(atomtype));

			if (PyObject_HasAttrString(self, "cutoff")) {
				PyObject* cutoff = PyObject_GetAttrString(self, "cutoff");
				pfunc->cutoff = PyFloat_AsDouble(cutoff);
				pfunc->buffer = pfunc->cutoff / 20;
				if (PyErr_Occurred() != NULL) {
					Py_DECREF(cutoff);
					return false;
				}
				Py_DECREF(cutoff);

				if (PyObject_HasAttrString(self, "buffer")) {
					PyObject* buffer = PyObject_GetAttrString(self, "buffer");
					pfunc->buffer = PyFloat_AsDouble(buffer);
					if (PyErr_Occurred() != NULL) {
						Py_DECREF(buffer);
						return false;
					}
					Py_DECREF(buffer);
				}

				if (PyObject_HasAttrString(model, "group_id")) {
					boost::tuple<PyObject*, Index*, size_t> gid = getIndexArray(model, "group_id");
					if (boost::get<0>(gid) == NULL) {
						return false;
					}
					pfunc->groupID = boost::get<1>(gid);
					Py_DECREF(boost::get<0>(gid));
				}
			}
		}

		if (hasCharge) {
			boost::tuple<PyObject*, RealNumber*, size_t> charges = getRealArray(model, "charges");
			if (boost::get<0>(charges) == NULL)  {
				return false;
			}
			else if (boost::get<2>(charges) != pfunc->nAtoms) {
				Py_DECREF(boost::get<0>(charges));
				return false;
			}
			pfunc->charges = boost::get<1>(charges);
			Py_DECREF(boost::get<0>(charges));

			PyObject* cc = PyObject_GetAttrString(self, "columnb_constant");
			assert (cc != NULL);
			Py_ssize_t cci = PyInt_AsSsize_t(cc);
			if (cci == -1) {
				return false;
			}
			pfunc->k_columnb = pfunc->pars + cci;
			Py_DECREF(cc);
		}
		else {
			pfunc->charges = pfunc->k_columnb = NULL;
		}

		return true;
	
	}


	void screenPairList(PyObject* filter, PyObject* fffunc, int rank) {
		
		static const char masks[] = {1, 2, 4, 8, 16, 32, 64, 128};
		char* pFilter = (char*)PyArray_DATA(filter);
		if (pFilter == NULL) {
			return;
		}
		boost::tuple<PyObject*, Index*, size_t> lists = getIndexArray(fffunc, "lists");
		if (boost::get<0>(lists) == NULL) {
			return;
		}
		Index* pList = boost::get<1>(lists);

		size_t n = boost::get<2>(lists);
		for (int  i = 0; i < n; i += rank+1) {
			Index a1 = pList[i+0];
			Index a2 = pList[i+rank-1];
			if (a1 < a2) {
				std::swap(a1, a2);
			}
			size_t offset = a1 * (a1 - 1) / 2 + a2;
			size_t j = offset / 8, k = offset % 8;
			pFilter[j] &= ~masks[k];
		}		
	}
} }


